#!/usr/bin/env python3
import os
import sys

CURDIR = os.path.dirname(os.path.realpath(__file__))
sys.path.insert(0, os.path.join(CURDIR, "helpers"))

CLICKHOUSE_URL = os.environ.get("CLICKHOUSE_URL")
CLICKHOUSE_TMP = os.environ.get("CLICKHOUSE_TMP")

from pure_http_client import ClickHouseClient

client = ClickHouseClient()


def run_test(data_format, gen_data_template, settings):
    print(data_format)
    client.query("TRUNCATE TABLE t_async_insert")

    expected = client.query(gen_data_template.format("TSV")).strip()
    data = client.query(
        gen_data_template.format(data_format), settings=settings, binary_result=True
    )

    insert_query = "INSERT INTO t_async_insert FORMAT {}".format(data_format)
    client.query_with_data(insert_query, data, settings=settings)

    result = client.query("SELECT * FROM t_async_insert FORMAT TSV").strip()
    if result != expected:
        print(
            "Failed for format {}.\nExpected:\n{}\nGot:\n{}\n".format(
                data_format, expected, result
            )
        )
        exit(1)


formats = (
    client.query(
        "SELECT name FROM system.formats WHERE is_input AND is_output \
    AND name NOT IN ('CapnProto', 'RawBLOB', 'Template', 'ProtobufSingle', 'LineAsString', 'Protobuf', 'ProtobufList', 'Npy') ORDER BY name"
    )
    .strip()
    .split("\n")
)

# Generic formats
client.query("DROP TABLE IF EXISTS t_async_insert")
client.query(
    "CREATE TABLE t_async_insert (id UInt64, s String, arr Array(UInt64)) ENGINE = Memory"
)
gen_data_query = "SELECT number AS id, toString(number) AS s, range(number) AS arr FROM numbers(10) FORMAT {}"

for data_format in formats:
    run_test(
        data_format,
        gen_data_query,
        settings={"async_insert": 1, "wait_for_async_insert": 1},
    )

# LineAsString
client.query("DROP TABLE IF EXISTS t_async_insert")
client.query("CREATE TABLE t_async_insert (s String) ENGINE = Memory")
gen_data_query = "SELECT toString(number) AS s FROM numbers(10) FORMAT {}"

run_test(
    "LineAsString",
    gen_data_query,
    settings={"async_insert": 1, "wait_for_async_insert": 1},
)

# TODO: add CapnProto and Protobuf

print("OK")
